/*
 * Copyright 2009 XueSong Guo.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cn.webwheel.utils.fileupload;

import cn.webwheel.Filter;
import cn.webwheel.FilterChain;
import cn.webwheel.WebContext;
import cn.webwheel.di.*;
import cn.webwheel.di.utils.TypeLiteral;
import cn.webwheel.utils.DIFilter;
import cn.webwheel.utils.WebParam;
import org.apache.commons.fileupload.FileItemIterator;
import org.apache.commons.fileupload.FileItemStream;
import org.apache.commons.fileupload.FileUploadException;
import org.apache.commons.fileupload.servlet.ServletFileUpload;
import org.apache.commons.fileupload.util.Streams;

import javax.servlet.http.HttpServletRequest;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class FileUploadFilter implements Filter {

    /**
     * request中存储文件上传信息的attribute名字
     */
    public static final String WFU = "webwheel.wfu";

    private int fileSizeMax;
    private int sizeMax;
    private String charset;
    private File tempDir;

    private Container container;
    private DIFilter diFilter;

    public FileUploadFilter(Container container, DIFilter diFilter, int fileSizeMax, int sizeMax, String charset, File tempDir) {
        this.container = container;
        this.diFilter = diFilter;
        this.fileSizeMax = fileSizeMax;
        this.sizeMax = sizeMax;
        this.charset = charset;
        this.tempDir = tempDir;
    }

    @SuppressWarnings("unchecked")
    public void init() throws Exception {
        container.bind(new Key<Map<String, List<FileInfo>>>(new TypeLiteral<Map<String, List<FileInfo>>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<Map<String, List<FileInfo>>>() {
            public Map<String, List<FileInfo>> get(Key<Map<String, List<FileInfo>>> mapKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                return (Map<String, List<FileInfo>>) request.getAttribute(WFU);
            }
        });
        container.bind(new Key<Map<String, List<File>>>(new TypeLiteral<Map<String, List<File>>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<Map<String, List<File>>>() {
            public Map<String, List<File>> get(Key<Map<String, List<File>>> mapKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if(map==null) return null;
                Map<String, List<File>> nmap = new HashMap<String, List<File>>(map.size());
                for (Map.Entry<String, List<FileInfo>> entry : map.entrySet()) {
                    List<File> list = new ArrayList<File>(entry.getValue().size());
                    for (FileInfo info : entry.getValue()) {
                        list.add(info.getFile());
                    }
                    nmap.put(entry.getKey(), list);
                }
                return nmap;
            }
        });

        container.bind(new Key<Map<String, FileInfo[]>>(new TypeLiteral<Map<String, FileInfo[]>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<Map<String, FileInfo[]>>() {
            public Map<String, FileInfo[]> get(Key<Map<String, FileInfo[]>> key, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                Map<String, FileInfo[]> nmap = new HashMap<String, FileInfo[]>(map.size());
                for (Map.Entry<String, List<FileInfo>> entry : map.entrySet()) {
                    nmap.put(entry.getKey(), entry.getValue().toArray(new FileInfo[entry.getValue().size()]));
                }
                return nmap;
            }
        });
        container.bind(new Key<Map<String, File[]>>(new TypeLiteral<Map<String, File[]>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<Map<String, File[]>>() {
            public Map<String, File[]> get(Key<Map<String, File[]>> key, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                Map<String, File[]> nmap = new HashMap<String, File[]>(map.size());
                for (Map.Entry<String, List<FileInfo>> entry : map.entrySet()) {
                    File[] list = new File[entry.getValue().size()];
                    for (int i = 0; i < list.length; i++) {
                        list[i] = entry.getValue().get(i).getFile();
                    }
                    nmap.put(entry.getKey(), list);
                }
                return nmap;
            }
        });

        container.bind(new Key<List<FileInfo>>(new TypeLiteral<List<FileInfo>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<List<FileInfo>>() {
            public List<FileInfo> get(Key<List<FileInfo>> listKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                return map.get(data);
            }
        });
        container.bind(new Key<List<File>>(new TypeLiteral<List<File>>(){}.getType(), WebParam.class.getSimpleName()), new Provider<List<File>>() {
            public List<File> get(Key<List<File>> listKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                List<FileInfo> list = map.get(data);
                if (list == null) return null;
                List<File> nlist = new ArrayList<File>(list.size());
                for (FileInfo info : list) {
                    nlist.add(info.getFile());
                }
                return nlist;
            }
        });

        container.bind(new Key<FileInfo[]>(FileInfo[].class, WebParam.class.getSimpleName()), new Provider<FileInfo[]>() {
            public FileInfo[] get(Key<FileInfo[]> key, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                List<FileInfo> list = map.get(data);
                if (list == null) return null;
                FileInfo[] infos = new FileInfo[list.size()];
                for (int i = 0; i < infos.length; i++) {
                    infos[i] = list.get(i);
                }
                return infos;
            }
        });
        container.bind(new Key<File[]>(File[].class, WebParam.class.getSimpleName()), new Provider<File[]>() {
            public File[] get(Key<File[]> key, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                List<FileInfo> list = map.get(data);
                if (list == null) return null;
                File[] nlist = new File[list.size()];
                for (int i = 0; i < nlist.length; i++) {
                    nlist[i] = list.get(i).getFile();
                }
                return nlist;
            }
        });

        container.bind(new Key<FileInfo>(FileInfo.class, WebParam.class.getSimpleName()), new Provider<FileInfo>() {
            public FileInfo get(Key<FileInfo> fileInfoKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                List<FileInfo> list = map.get(data);
                if (list == null) return null;
                return list.get(0);
            }
        });
        container.bind(new Key<File>(File.class, WebParam.class.getSimpleName()), new Provider<File>() {
            public File get(Key<File> fileKey, String data, NativeProvider<?> nativeProvider) {
                HttpServletRequest request = container.getInstance(new Key<HttpServletRequest>(HttpServletRequest.class, Desire.DefaultName), Desire.DefaultData);
                if(request==null) return null;
                Map<String, List<FileInfo>> map = (Map<String, List<FileInfo>>) request.getAttribute(WFU);
                if (map == null) return null;
                List<FileInfo> list = map.get(data);
                if (list == null) return null;
                return list.get(0).getFile();
            }
        });
    }

    @SuppressWarnings("ResultOfMethodCallIgnored")
    public boolean process(FilterChain chain) throws Exception {

        HttpServletRequest request = chain.getWebContext().getRequest();

        if (!ServletFileUpload.isMultipartContent(request)) return chain.go();

        ServletFileUpload fileUpload = new ServletFileUpload();

        int sizeMax = this.sizeMax;
        String charset = this.charset;
        int fileSizeMax = this.fileSizeMax;

        FileUploadSettings settings = chain.getActionClass().getAnnotation(FileUploadSettings.class);
        if (settings != null) {
            sizeMax = settings.sizeMax();
            charset = settings.charset();
            fileSizeMax = settings.fileSizeMax();
        }
        if (sizeMax > 0) {
            fileUpload.setSizeMax(sizeMax);
        }
        if(charset!=null) {
            fileUpload.setHeaderEncoding(charset);
        }
        if (fileSizeMax > 0) {
            fileUpload.setFileSizeMax(fileSizeMax);
        }

        Map<String, List<String>> params = new HashMap<String, List<String>>();
        ThreadLocal<WebContext> threadLocal = diFilter.getThreadLocal();
        WebContext ctx = threadLocal.get();
        threadLocal.set(new WebContextWrapper(ctx, params));

        Map<String, List<FileInfo>> fileInfoMap = new HashMap<String, List<FileInfo>>();
        try {
            FileItemIterator it;

            try {
                it = fileUpload.getItemIterator(request);
            } catch (IOException e) {
                // 忽略上传中断错误
                return true;
            } catch (FileUploadException e) {
                if (FileUploadExceptionAware.class.isAssignableFrom(chain.getActionClass())) {
                    ((FileUploadExceptionAware)chain.getAction()).occur(e);
                    return chain.go();
                } else {
                    throw e;
                }
            }

            while (true) {
                FileItemStream item;
                try {
                    if (!it.hasNext()) break;
                    item = it.next();
                } catch (IOException e) {
                    // 忽略上传中断错误
                    return true;
                } catch (FileUploadException e) {
                    if (FileUploadExceptionAware.class.isAssignableFrom(chain.getActionClass())) {
                        ((FileUploadExceptionAware)chain.getAction()).occur(e);
                        break;
                    } else {
                        throw e;
                    }
                }
                String name = item.getFieldName();
                if (item.isFormField()) {
                    List<String> list = params.get(name);
                    if (list == null) params.put(name, list = new ArrayList<String>(1));
                    String s;
                    try {
                        s = Streams.asString(item.openStream(), charset == null ? "utf-8" : charset);
                    } catch (IOException e) {
                        // 忽略上传中断错误
                        return true;
                    }
                    list.add(s);
                } else if (!item.getName().isEmpty()) {
                    InputStream is;
                    try {
                        is = item.openStream();
                    } catch (IOException e) {
                        // 忽略上传错误
                        return true;
                    }

                    File file = File.createTempFile("wfu", null, tempDir);
                    file.deleteOnExit();

                    List<FileInfo> list = fileInfoMap.get(name);
                    if (list == null) fileInfoMap.put(name, list = new ArrayList<FileInfo>(1));
                    FileInfoImpl fileInfo = new FileInfoImpl(file);
                    list.add(fileInfo);

                    FileOutputStream fos = new FileOutputStream(file);
                    byte[] buf = new byte[8192];
                    try {
                        while (true) {
                            int rd;
                            try {
                                rd = is.read(buf);
                                if (rd == -1) break;
                            } catch (IOException e) {
                                // 忽略上传错误
                                return true;
                            }
                            fos.write(buf, 0, rd);
                        }
                    } finally {
                        try {
                            fos.close();
                        } catch (IOException e) {/*ignore*/}
                    }

                    fileInfo.fileName = item.getName();
                    fileInfo.contentType = item.getContentType();
                    /* not implement
                    item.getHeaders();
                    */
                }
            }

            request.setAttribute(WFU, fileInfoMap);

            return chain.go();
        } finally {
            threadLocal.set(ctx);
            for (Map.Entry<String, List<FileInfo>> entry : fileInfoMap.entrySet()) {
                for (FileInfo info : entry.getValue()) {
                    info.getFile().delete();
                }
            }
        }
    }

    public void destroy() {
    }

    private static class FileInfoImpl implements FileInfo {

        File file;
        String fileName;
        String contentType;

        public FileInfoImpl(File file) {
            this.file = file;
        }

        public File getFile() {
            return file;
        }

        public String getFileName() {
            return fileName;
        }

        public String getContentType() {
            return contentType;
        }

        public String getSimpleFileName() {
            return fileName.substring(Math.max(fileName.lastIndexOf('/'), fileName.lastIndexOf('\\')) + 1);
        }

        public String getExtension() {
            String s = getSimpleFileName();
            int i = s.lastIndexOf('.');
            if (i == -1) return "";
            return s.substring(i);
        }
    }
}
